import cv2
import os
import numpy as np
from ultralytics import YOLO

# 配置参数
input_folder = "D:/Desk/test1/content/"
output_folder = "D:/Desk/test1/result/"
target_ratio = 3/4  # 严格宽高比 3:4
model_path = "../yolov8x-cls.pt"

# 安全加载模型
try:
    model = YOLO(model_path) if os.path.exists(model_path) else None
except:
    model = None

def get_vertical_center(img):
    """智能获取垂直中心点（优先衣物位置）"""
    h, w = img.shape[:2]
    
    # 优先级1：YOLO衣物检测
    if model is not None:
        try:
            results = model.predict(img, verbose=False)
            for r in results:
                if r.boxes is not None:
                    boxes = r.boxes.xyxy.cpu().numpy()
                    if len(boxes) > 0:
                        best = boxes[np.argmax(boxes[:, 4])]  # 取最高置信度框
                        return int((best[1] + best[3]) / 2)  # 返回垂直中心
        except:
            pass
    
    # 优先级2：人体检测（修复判断逻辑）
    hog = cv2.HOGDescriptor()
    hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())
    regions, _ = hog.detectMultiScale(img)
    if regions is not None and len(regions) > 0:
        _, y, _, h = max(regions, key=lambda r: r[2]*r[3])
        return y + h//2
    
    # 保底：垂直居中
    return h//2

def strict_ratio_crop(img):
    """严格比例裁剪核心逻辑（修复水平中心检测）"""
    h, w = img.shape[:2]
    current_ratio = w / h
    
    # 比例恰好为3:4时直接返回
    if abs(current_ratio - target_ratio) < 0.01:
        return img
    
    if current_ratio < target_ratio:
        # 情况1：原图"较窄" → 保持原宽度，计算目标高度
        target_h = int(w / target_ratio)
        
        # 获取垂直中心点
        center_y = get_vertical_center(img)
        
        # 计算裁剪范围
        y_start = max(0, center_y - target_h//2)
        y_end = y_start + target_h
        
        # 边界溢出处理
        if y_end > h:
            y_start = max(0, h - target_h)
            y_end = h
        
        return img[y_start:y_end, 0:w]
    else:
        # 情况2：原图"较宽" → 保持原高度，计算目标宽度
        target_w = int(h * target_ratio)
        
        # 获取水平中心点（修复逻辑）
        cx = get_horizontal_center(img)
        
        # 计算裁剪范围
        x_start = max(0, cx - target_w//2)
        x_end = x_start + target_w
        
        # 边界溢出处理
        if x_end > w:
            x_start = max(0, w - target_w)
            x_end = w
        
        return img[0:h, x_start:x_end]

def get_horizontal_center(img):
    """修复后的水平中心检测"""
    h, w = img.shape[:2]
    
    # 优先级1：YOLO衣物检测
    if model is not None:
        try:
            results = model.predict(img, verbose=False)
            for r in results:
                if r.boxes is not None:
                    boxes = r.boxes.xyxy.cpu().numpy()
                    if len(boxes) > 0:
                        best = boxes[np.argmax(boxes[:, 4])]  # 取最高置信度框
                        return int((best[0] + best[2]) / 2)
        except:
            pass
    
    # 优先级2：人体检测（修复判断逻辑）
    hog = cv2.HOGDescriptor()
    hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())
    regions, _ = hog.detectMultiScale(img)
    if regions is not None and len(regions) > 0:
        x, _, w, _ = max(regions, key=lambda r: r[2]*r[3])
        return x + w//2
    
    # 保底：水平居中
    return w//2

# 批量处理
os.makedirs(output_folder, exist_ok=True)
for filename in os.listdir(input_folder):
    if not filename.lower().endswith(('.png', '.jpg', '.jpeg')):
        continue

    img_path = os.path.join(input_folder, filename)
    img = cv2.imread(img_path)
    if img is None:
        continue

    # 执行严格比例裁剪
    cropped = strict_ratio_crop(img)
    
    # 强制验证输出比例
    h, w = cropped.shape[:2]
    assert abs(w/h - target_ratio) < 0.02, f"比例错误: {w}/{h} = {w/h:.3f}"
    
    cv2.imwrite(os.path.join(output_folder, filename), cropped)

print("处理完成！输出比例严格保持3:4")